import base64
import datetime
import logging
import os

from paddleocr import PaddleOCR
import re
import time
import torch
from torchvision import transforms
from PIL import Image

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 加密
aes_key = b'L03-Cumminschine'
iv_key = b'Cty-china-ER-app'

# 图片路径
image_file_path = "C:/Users/qk965/Desktop/ocr_local_test/images/"
model_path = r"../models/model_draft_1129.pth"
# loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
ocr = PaddleOCR(
    use_angle_cls=True,
    cls_model_dir=r"../models/cls",
    det_model_dir=r"../models/det",
    rec_model_dir=r"../models/rec"
)
# 角度不对 模糊 假的  未识别
class_labels = ['您拍摄照片的角度倾斜', '您拍摄的照片比较模糊', '您拍摄的照片可能存在翻拍', '']


# 上传文件
# type (1  ESN) (2 OEMName、ChassisNumber、VIN) (3 OEMName、VIN) (4、WarrantyStartDate、ChassisNumber) (5、WarrantyStartDate)
def upload(file_content, file_name, ocr_type):
    file_path = image_file_path + file_name
    logger.info(f"saving file into : {file_path}")
    __save_file(file_path, file_content)

    ocr_result = __run_model(file_path, ocr_type)
    msg = ""
    if ocr_result is None:
        msg = __single_cf_test(file_path)

    logger.info("识别内容" + str(ocr_result))

    __clear_file(file_path)

    return {"code": 1, "data": ocr_result, 'msg': msg}


def __run_model(img_path, ocr_type):
    result = ocr.ocr(img_path, cls=True)
    # logger.info("ocr 结果解析前：" + str(result))
    if result and len(result) > 0:
        result_arr = [line[1][0] for line in result[0]]
        logger.info("图片文字：" + str(result_arr))
        return __resolve_ocr_result(result_arr, ocr_type)
    else:
        logger.warning("解析失败")
        return None


def get_vin(ocr_result, vin_count, vin_list):
    v_vin = re.findall("[L]{1}[A-Z0-9]{16}", ocr_result)
    if (v_vin != []) & (vin_count == 0):
        v_vin[0] = v_vin[0].replace('O', '0')
        vin_list.append(v_vin[0])
        vin_count += 1
    return vin_count


def get_esn(ocr_result, esn_count, esn_list):
    v_esn = re.findall("[7]{1}[0-9]{7}", ocr_result)
    if (v_esn != []) & (esn_count == 0):
        esn_list.append(v_esn[0])
        esn_count += 1
    return esn_count


def get_date(ocr_result, date_count, date_list):
    v_date = re.findall("\d{4}-\d{2}-\d{2}", ocr_result)
    if (v_date != []) & (date_count == 0):
        date_list.append(v_date[0])
        date_count += 1
    if (v_date != []) & (date_count != 0):
        try:
            v_date[0] = v_date[0].replace('88', '08')
            date_list[-1] = date_list[-1].replace('88', '08')
            date_a = datetime.strptime(v_date[0], '%Y-%m-%d')
            date_b = datetime.strptime(date_list[-1], '%Y-%m-%d')
            if date_a < date_b:
                date_list[-1] = v_date[0]
        except:
            pass
    return date_count


def get_model_on(ocr_result, model_on_count, model_on_list):
    v_model_rd = re.findall(r'[\u4e00-\u9fa5]+牌[A-Z0-9]+.*', ocr_result)
    if (v_model_rd != []) & (model_on_count == 0):
        # print('model_rd',v_model_rd[0])
        v_model_rd[0] = v_model_rd[0].replace('品牌型号', '')
        model_on_list.append(v_model_rd[0])
        model_on_count += 1
    return model_on_count


def get_model_off(ocr_result, model_off_count, model_off_list):
    v_model_off = re.findall(r'^F[0-9]+\.[A-Z0-9]+$', ocr_result)
    if (v_model_off != []) & (model_off_count == 0):
        model_off_list.append(v_model_off[0])
        model_off_count += 1

    v_model_off = re.findall(r'[X](?=.*[0-9])[0-9a-zA-Z]{5,10}(?<![kK][wW]|kg)$', ocr_result)
    if (v_model_off != []) & (model_off_count == 0):
        model_off_list.append(v_model_off[0])
        model_off_count += 1
    return model_off_count


def get_company(ocr_result, company_count, company_list):
    v_comp = re.findall(r'.*有限.*', ocr_result)
    if (v_comp != []) & (company_count == 0):
        company_list.append(v_comp[0])
        company_count += 1
    return company_count


def __resolve_ocr_result(result_arr, ocr_type):
    vins = []
    esns = []
    dates = []
    model_on = []  # 道路
    model_off = []  # 非道路
    comp = []

    vin_count = 0
    esn_count = 0
    date_count = 0
    model_on_count = 0
    model_off_count = 0
    company_count = 0
    for i in range(len(result_arr)):
        result_value = str(result_arr[i])
        vin_count = get_vin(result_value, vin_count, vins)
        esn_count = get_esn(result_value, esn_count, esns)
        date_count = get_date(result_value, date_count, dates)
        model_on_count = get_model_on(result_value, model_on_count, model_on)
        model_off_count = get_model_off(result_value, model_off_count, model_off)
        company_count = get_company(result_value, company_count, comp)

    ocr_results = {}
    if vins and ocr_type in [2, 3]:
        ocr_results["VIN"] = vins[0]

    if esns and ocr_type == 1:
        ocr_results["ESN"] = esns[0]

    if dates and ocr_type in [4, 5]:
        ocr_results["WarrantyStartDate"] = dates[0]

    if model_on and ocr_type == 4:
        ocr_results["ChassisNumber"] = model_on[0]

    if model_off and ocr_type == 2:
        ocr_results["ChassisNumber"] = model_off[0]

    if comp and ocr_type in [2, 3]:
        ocr_results["OEMName"] = comp[0]

    return ocr_results


# 识别失败-报错处理
def __single_cf_test(image_path):
    global loaded_model
    image = Image.open(image_path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0)
    if loaded_model is None:
        loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
    loaded_model.eval()
    with torch.no_grad():
        output = loaded_model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    predicted_class = torch.argmax(probabilities).item()
    predicted_label = class_labels[predicted_class]
    logger.info(f'--- \n Predicted Class: {predicted_label}')
    print(f'Class Probability: {probabilities[predicted_class].item()}')
    return predicted_label
    # classification result========================


# save file to local path
def __save_file(file_path, file_content):
    try:
        with open(file_path, "wb") as lf:
            lf.write(file_content)
        return True
    except Exception as e:
        logger.error('成功保存失败发生错误：', str(e))
        return False

def __clear_file(file_path):
    try:
        os.remove(file_path)
    except Exception as e:
        logger.error('clear file error', str(e))

if __name__ == '__main__':
    event = {
        "headers": {
            "type": "2",
            "filename": "test.png"
        },
        # base64 encoded image
        "body": "123"
    }

    headers = event.get("headers")
    body = event.get('body')
    ocr_type = int(headers.get("type"))
    file_name = headers.get("filename")

    if type is None:
        raise ValueError("type is invalid")
    if body is None:
        raise ValueError("body is invalid")
    if file_name is None:
        raise ValueError("filename is invalid")

    with open("C:/Users/qk965/Desktop/ocr_local_test/local/test_1.png", "rb") as img:
        img_str = base64.b64encode(img.read())

    print(f"img_str: {img_str}")
    file_content = base64.b64decode(img_str)

    result = upload(file_content, file_name, ocr_type)
    print(f"result: {result}")
